package com.example.spring_ai.a2a.client.service;

import com.example.spring_ai.a2a.server.bean.TaskRequest;
import com.example.spring_ai.a2a.server.bean.TaskResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.http.*;
import org.springframework.stereotype.Service;
import org.springframework.web.client.RestTemplate;
import org.springframework.web.reactive.function.client.WebClient;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import reactor.core.publisher.Flux;
import reactor.util.retry.Retry;

import java.io.IOException;
import java.time.Duration;
import java.util.Map;
import java.util.UUID;

@Service
@Slf4j
public class AgentClientService {
    private final RestTemplate restTemplate;
    private final WebClient webClient;

    String baseUrl = "http://localhost:8181";

    public AgentClientService(RestTemplate restTemplate) {
        this.restTemplate = restTemplate;
        this.webClient = WebClient.builder()
                .baseUrl(baseUrl)
                .defaultHeader(HttpHeaders.CONTENT_TYPE, MediaType.APPLICATION_JSON_VALUE)
                .build();
    }

    public SseEmitter tasksPlan(String tasks) throws InterruptedException {

        // 获取WeatherAgent的AgentCard
        log.info("[client] Start to get AgentCard");
        Map<String, Object> agentCard = restTemplate.getForObject(
                "http://localhost:8181/.well-known/agent.json",
                Map.class
        );
        if (agentCard == null) {
            throw new RuntimeException("Failed to fetch agent card");
        }
        log.info("[client] Get AgentCard: %s".formatted(agentCard));

        // 构造A2A标准任务请求
        Map<String, Object> endpoints = (Map<String, Object>) agentCard.get("endpoints");
        String taskSubmitUrl = baseUrl + endpoints.get("task_send");
        HttpHeaders headers = new HttpHeaders();
        headers.set("Authorization", "SECRET_KEY");
        String taskId = UUID.randomUUID().toString();
        Map<String, String> params = Map.of("tasks", tasks);
        ResponseEntity<TaskResponse> response = restTemplate.exchange(
                taskSubmitUrl,
                HttpMethod.POST,
                new HttpEntity<>(new TaskRequest(taskId, params), headers),
                TaskResponse.class
        );
        log.info("[client] Step task_send end, {}", response);

        // SSE订阅
        SseEmitter sseEmitter = new SseEmitter(60_000L);
        Flux<TaskResponse> eventStream = webClient.get()
                .uri(uriBuilder -> uriBuilder
                        .path(endpoints.get("task_get").toString())
                        .queryParam("taskId", taskId)
                        .build())
                .header("Authorization", "SECRET_KEY")
                .accept(MediaType.TEXT_EVENT_STREAM)
                .retrieve()
                .bodyToFlux(TaskResponse.class)
                .filter(resp -> taskId.equals(resp.getTaskId()))
                .retryWhen(Retry.backoff(3, Duration.ofSeconds(1)));

        eventStream.subscribe(
                resp -> {
                    try {
                        sseEmitter.send(resp);
                        if ("completed".equals(resp.getStatus())) {
                            sseEmitter.complete(); // 服务端标记完成时关闭
                        }
                    } catch (IOException e) {
                        throw new RuntimeException(e);
                    }
                },
                error -> log.info("SSE连接异常: " + error.getMessage())
        );
        
        log.info("[client] Step sse_subscribe end ");

        return sseEmitter;
    }
}
